import json
from decimal import ROUND_HALF_UP, Decimal

from agentops import logger
from agentops.exceptions import ApiServerException
from agentops.helpers import get_ISO_time, filter_unjsonable
from agentops.http_client import HttpClient
from typing import Union
from datetime import datetime


def get_analytics(session) -> dict[str, Union[Decimal, str]]:
    if not session.end_timestamp:
        session.end_timestamp = get_ISO_time()

    formatted_duration = _format_duration(
        session.init_timestamp, session.end_timestamp
    )
    token_cost = _get_token_cost(session)
    if token_cost == "unknown" or token_cost is None:
        token_cost_d = Decimal(0)
    else:
        token_cost_d = Decimal(token_cost)
    formatted_cost = _format_token_cost(token_cost_d)

    return {
        "llm_calls": session.event_counts["llms"],
        "tool_calls": session.event_counts["tools"],
        "actions": session.event_counts["actions"],
        "errors": session.event_counts["errors"],
        "duration": formatted_duration,
        "cost": formatted_cost,
    }


def _format_duration(start_time, end_time):
    start = datetime.fromisoformat(start_time.replace("Z", "+00:00"))
    end = datetime.fromisoformat(end_time.replace("Z", "+00:00"))
    duration = end - start

    hours, remainder = divmod(duration.total_seconds(), 3600)
    minutes, seconds = divmod(remainder, 60)

    parts = []
    if hours > 0:
        parts.append(f"{int(hours)}h")
    if minutes > 0:
        parts.append(f"{int(minutes)}m")
    parts.append(f"{seconds:.1f}s")

    return " ".join(parts)


def _get_token_cost(session):
    with session.lock:
        payload = {"session": session.__dict__}
        try:
            res = HttpClient.post(
                f"{session.config.endpoint}/v2/update_session",
                json.dumps(filter_unjsonable(payload)).encode("utf-8"),
                jwt=session.jwt,
            )
        except ApiServerException as e:
            return logger.error(f"Could not end session - {e}")

    logger.debug(res.body)
    return res.body.get("token_cost", "unknown")


def _format_token_cost(token_cost_d):
    return (
        "{:.2f}".format(token_cost_d)
        if token_cost_d == 0
        else "{:.6f}".format(
            token_cost_d.quantize(Decimal("0.000001"), rounding=ROUND_HALF_UP)
        )
    )